"""
百度文心大模型客户端
"""

import os
from typing import Any, Dict, Generator, List

import qianfan

from ..decorators.performance import timeit
from ..logger import setup_logger
from .base_client import BaseClient, ChatResponse

# 设置日志记录器
logger = setup_logger(__name__)


class WenXinClientBase(BaseClient):
    """文心大模型基础客户端"""

    def __init__(self, model_name):
        super().__init__(model_name)
        self._setup_client()

    def _setup_client(self):
        """初始化客户端，由子类实现"""
        pass

    def _handle_error(self, e: Exception) -> ChatResponse:
        """统一错误处理"""
        logger.error(f"Chat error with {self.model_name}: {e}")
        return ChatResponse(content="", status=False, error_msg=str(e))


class WenXinClientV1(WenXinClientBase):
    """
    文心大模型V1版本客户端
    参考文档： https://cloud.baidu.com/doc/WENXINWORKSHOP/s/7lq3ft3pb
    """

    def _setup_client(self):
        os.environ["QIANFAN_ACCESS_KEY"] = self.model_conf["ak"]
        os.environ["QIANFAN_SECRET_KEY"] = self.model_conf["sk"]
        self.client = qianfan.ChatCompletion()

    def chat(self, messages) -> ChatResponse:
        try:
            resp = self.client.do(model=self.model_name,
                                  messages=messages,
                                  stream=False)
            return ChatResponse(content=resp["body"]["result"])
        except Exception as e:
            return self._handle_error(e)

    def chat_stream(self, messages) -> Generator[str, None, None]:
        try:
            # 虽然V1版本不支持流式输出，但我们可以模拟一个流式输出
            resp = self.chat(messages)
            if resp.status:
                yield resp.content
            else:
                yield f"Error: {resp.error_msg}"
        except Exception as e:
            logger.error(f"Stream chat error with {self.model_name}: {e}")
            yield f"Error: {str(e)}"


class WenXinClientV2(WenXinClientBase):
    """文心大模型V2版本客户端"""

    def _setup_client(self):
        self.client = qianfan.Qianfan(
            access_key=self.model_conf["ak"],
            secret_key=self.model_conf["sk"],
        )

    @timeit
    def chat(self, messages) -> ChatResponse:
        try:
            completion = self.client.chat.completions.create(
                model=self.model_conf["model_name"].lower(),
                messages=messages,
                stream=False,
            )
            result = completion.choices[0].message.content
            logger.info(f"{self.model_name} response: {result}")
            return ChatResponse(content=result)
        except Exception as e:
            return self._handle_error(e)

    @timeit
    def chat_stream(self, messages) -> Generator[str, None, None]:
        try:
            completion = self.client.chat.completions.create(
                model=self.model_conf["model_name"],
                messages=messages,
                stream=True,
            )
            for chunk in completion:
                if chunk.choices[0].delta.content is not None:
                    yield chunk.choices[0].delta.content
        except Exception as e:
            logger.error(f"Stream chat error with {self.model_name}: {e}")
            yield f"Error: {str(e)}"
